# -*- coding: utf-8 -*-
"""
   File Name：     inference
   Description :  frcnn预测
   Author :       mick.yi
   date：          2019/2/13
"""
import argparse
import os
import sys

import matplotlib
import numpy as np

matplotlib.use('Agg')
from matplotlib import pyplot as plt
from faster_rcnn.preprocess.input import VocDataset
from faster_rcnn.utils import image as image_utils
from faster_rcnn.utils import visualize, np_utils
from faster_rcnn.config import current_config as config
from pb_inference import pb_inference


def class_map_to_id_map(class_mapping):
    id_map = {}
    for k, v in class_mapping.items():
        id_map[v] = k
    return id_map


def main(args):
    os.environ["CUDA_VISIBLE_DEVICES"] = config.INFERENCE_GPU_ID
    # 覆盖参数
    config.IMAGES_PER_GPU = 1
    config.GPU_COUNT = 1
    # 加载数据
    dataset = VocDataset(config.voc_path, class_mapping=config.CLASS_MAPPING)
    dataset.prepare()
    all_img_info = [info for info in dataset.get_image_info_list() if info['type'] == 'test']  # 测试集

    # class map 转为 id map
    id_mapping = class_map_to_id_map(config.CLASS_MAPPING)

    def _show_inference(id, ax=None):
        image = image_utils.load_image(all_img_info[id]['filepath'])
        print("image_path: {}".format(all_img_info[id]['filepath']))
        image, image_meta, _ = image_utils.resize_image_and_gt(image,
                                                               config.IMAGE_MAX_DIM,
                                                               all_img_info[id]['boxes'])
        # print("image.shape: {}".format(image.shape))
        # print("image_meta: {}".format(image_meta.shape))
        result = pb_inference.pb_http_request({"input_image": image.tolist(),
                                               "input_image_meta": image_meta.tolist()
                                               })
        # print("result: {}".format(result))
        if not result.__contains__("predictions"):
            print("predict error: {}".format(result))
            return

        predict_result = result.get("predictions")[0]

        # --------- tf.save_model -----------
        # boxes = np.asarray(predict_result.get("detect_boxes"))
        # scores = np.asarray(predict_result.get("class_scores"))
        # class_ids = np.asarray(predict_result.get("detect_class_ids"))
        # detect_class_logits = np.asarray(predict_result.get("detect_class_logits"))
        # image_metas = np.asarray(predict_result.get("image_meta"))

        # ---------- tf.contrib.save_model ----------
        boxes = np.asarray(predict_result.get("lambda_2"))
        scores = np.asarray(predict_result.get("lambda_3"))
        class_ids = np.asarray(predict_result.get("lambda_4"))
        image_metas = np.asarray(predict_result.get("lambda_5"))
        detect_class_logits = np.asarray(predict_result.get("lambda_6"))


        boxes = np_utils.remove_pad(boxes)
        scores = np_utils.remove_pad(scores)[:, 0]
        class_ids = np_utils.remove_pad(class_ids)[:, 0]
        if len(class_ids) >= 0:
            print("boxes: {}".format(boxes))
            print("scores: {}".format(scores))
            print("class_ids: {}".format(class_ids))
            # print("detect_class_logits: {}".format(detect_class_logits))
            # print("image_metas: {}".format(image_metas))
        visualize.display_instances(image, boxes[:5],
                                    class_ids[:5],
                                    id_mapping,
                                    scores=scores[:5],
                                    ax=ax)
        print("boxes num:{}".format(boxes.shape[0]))

    # 随机展示9张图像

    # image_ids = np.random.choice(len(all_img_info), 1, replace=False)
    # image_ids = [1, 5, 6, 9, 10, 15, 16, 20, 22, 26]
    image_ids = [10, 19]
    fig = plt.figure(figsize=(20, 20))
    for idx, image_id in enumerate(image_ids):
        try:
            _show_inference(image_id, None)
        except Exception as err:
            print("exception: {}".format(err))
            continue
    fig.savefig('demo_images/inferece_examples_remote.{}.png'.format(image_ids[0]))


if __name__ == '__main__':
    parse = argparse.ArgumentParser()
    parse.add_argument("--weight_path", type=str, default=None, help="weight path")
    argments = parse.parse_args(sys.argv[1:])
    main(argments)
